from collections import Counter
from pathlib import Path
import pandas as pd
import reader
import os

CONFIG_FILENAME = 'config.ini'
DEMAND_FILENAME = 'demand.csv'
QOS_FILENAME = 'qos.csv'
SITE_BANDWIDTH_FILE = 'site_bandwidth.csv'

OUTPUT_FILE_NAME = 'solution.txt'

def get_abs_path(a: str, b: str):
    return Path.resolve(Path(a, b))

def simulator(data_base_path: str, output_base_path: str, verbose=False, error_check=False):
    qos_cons, V = reader.get_config_from(get_abs_path(data_base_path, CONFIG_FILENAME))
    
    demands, qos, site_bandwidth = [reader.get_data_from(get_abs_path(data_base_path, filename)) for filename in (DEMAND_FILENAME, QOS_FILENAME, SITE_BANDWIDTH_FILE)]
    
    ans = reader.get_ans_from(get_abs_path(output_base_path, OUTPUT_FILE_NAME))
    
    # demands = demands.drop(columns='mtime')
    qos = qos.set_index('site_name')
    site_bandwidth = site_bandwidth.set_index('site_name')
    bandwidth_seq = pd.DataFrame(columns=site_bandwidth.index)
    n_user = len(demands.columns) - 2

    sites = list(site_bandwidth.index)
    for index, ((_, _demand), _plan) in enumerate(zip(demands.groupby('mtime'), list(zip(*[iter(ans)]*n_user)))):
        plan = dict(_plan)
        demand = _demand.set_index('stream_id')
        demand_sum = demand.sum()

        if error_check:
            rec = {site: 0 for site in sites}
            for user, strategy in plan.items():
                assert (demand_sum[user] == 0 and strategy is None) or (demand_sum[user] > 0 and strategy is not None), f'user demand not matched, user {user} has demand {demand.loc[:, user]} at time {index} (start at 0), but got {strategy}'
                
                streams = []
                if strategy is not None:
                    for site, stream_ids in strategy.items():
                        assert qos.loc[site, user] < qos_cons, f'QOS constraint false, site {site} to user {user} at time {index} (start at 0)'
                        streams += list(stream_ids)
                        rec[site] += sum([demand.loc[stream_id, user] for stream_id in stream_ids])
                else:
                    continue

                observed_stream = Counter(streams)

                assert (pd.array(list(observed_stream.values())) == 1).all() and len(streams) == len(demand), f'stream can\' be splitted. user {user} at time {index} (start at 0) has illegal stream strategy: {[stream_id for stream_id in observed_stream if observed_stream[stream_id] != 1]}. strategy: {strategy}'

            assert (pd.Series(rec) <= site_bandwidth['bandwidth']).all(), f'site bandwidth exceeded at time {index} (start at 0), bandwidth: {rec}, real bandwidth: {site_bandwidth.loc[:, "bandwidth"]}'
        
        else:
            rec = {site: 0 for site in sites}

            for user, strategy in plan.items():
                if strategy is not None:
                    for site, stream_ids in strategy.items():
                        rec[site] += sum([demand.loc[stream_id, user] for stream_id in stream_ids])

        bandwidth_seq = pd.concat([bandwidth_seq, pd.DataFrame([rec])], axis=0, ignore_index=True)

    temp_seq = pd.concat([bandwidth_seq, pd.DataFrame([{k: -1 for k in bandwidth_seq.columns}])], axis=0, ignore_index=True)

    quantile = temp_seq.quantile(0.95, numeric_only=False, interpolation='higher')

    factor = (bandwidth_seq.sum() > 0).astype(int)

    quantile[quantile <= V] = V

    quantile[quantile > V] += (quantile[quantile > V] - V)**2 / site_bandwidth.loc[:, 'bandwidth'][quantile > V]

    quantile *= factor

    ans = sum(quantile)

    if verbose:
        output_dir = '__temp'

        os.makedirs(output_dir, exist_ok=True)

        print(f'costs: \n {quantile} \n ======= \n')

        print(f"quantile: {temp_seq.quantile(0.95, numeric_only=False, interpolation='higher')} \n")

        temp_seq.quantile(0.95, numeric_only=False, interpolation='higher').to_csv(f'{output_dir}/bandwidth_quant.csv')

        bandwidth_seq.to_csv(f'{output_dir}/bandwidth_seq.csv')

        quantile.to_csv(f"{output_dir}/costs.csv")

        print(f'==== \ndetailed information is saved in {output_dir}/bandwidth_quant.csv')

    print('ans: ', sum(temp_seq.quantile(0.95, numeric_only=False, interpolation='higher')))


if __name__=='__main__':
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('--data', type=str, default='../data', help='data path')
    parser.add_argument('--ans', type=str, default='../output', help='output file path')
    parser.add_argument('--error_check', action='store_true')
    parser.add_argument('--verbose', action='store_true')

    args = parser.parse_args()

    if args.verbose:
        print(args._get_kwargs())
    
    # print(simulator(args.data, args.ans))
    simulator(args.data, args.ans, args.verbose, args.error_check)